import cv2
import numpy as np
import os
import cv2
import numpy as np
from skimage.measure import label, regionprops
from tqdm import tqdm
import pandas as pd

# OCT厚层分析算法

"""
计算宽度
"""

def extract_center_line(binary_image):
    # 使用传统的细化算法进行中心线提取
    element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    done = False
    size = np.size(binary_image)
    skel = np.zeros(binary_image.shape, np.uint8)

    while not done:
        eroded = cv2.erode(binary_image, element)
        temp = cv2.dilate(eroded, element)
        temp = cv2.subtract(binary_image, temp)
        skel = cv2.bitwise_or(skel, temp)
        binary_image = eroded.copy()

        zeros = size - cv2.countNonZero(binary_image)
        if zeros == size:
            done = True

    return skel

def calculate_average_width(center_line_points, contours):
    average_widths = []

    for point in center_line_points:
        distances = []
        for contour in contours:
            distance = cv2.pointPolygonTest(contour, (int(point[1]), int(point[0])), True)
            distances.append(abs(distance))

        # 计算中心线到轮廓的最小距离
        min_distance = min(distances)

        # 将最小距离乘以2，即为连通域在该点的宽度
        width_at_point = min_distance * 2
        average_widths.append(width_at_point)

    # 计算平均宽度
    average_width = np.mean(average_widths)
    return average_width


def process_image(image_path, output_path_csv, color, slice):
    # 读取图像并进行二值化处理
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    _, binary_image = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)

    # 提取中心线
    center_line = extract_center_line(binary_image)
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 计算平均宽度
    center_line_points = np.column_stack(np.where(center_line > 0))
    average_width = calculate_average_width(center_line_points, contours)

    # 初始化layer变量为默认值
    layer = 'Unknown Layer'

    # 根据颜色值设置layer
    if all(color == [0, 0, 255]):
        layer = 'ILM-IPL'
    elif all(color == [0, 255, 0]):
        layer = 'IPL-OPL'
    elif all(color == [0, 255, 192]):
        layer = 'OPL-ISOS'
    elif all(color == [156, 255, 0]):
        layer = 'ISOS-RPE'
    elif all(color == [255, 0, 0]):
        layer = 'RPE-BM'

    data = {'Slice': [slice], 'Layer': [layer], 'Average Height (px)': [average_width]}
    df = pd.DataFrame(data)

    # 获取输出文件夹名和csv文件名
    folder_name = output_path_csv.split('/')[-1]
    csv_file_name = f"{folder_name}.csv"

    # 构建完整的输出CSV路径
    full_output_path = f"{output_path_csv}/{csv_file_name}"

    df.to_csv(full_output_path, mode='a', index=False, header=(not os.path.isfile(full_output_path)))

# 其余代码保持不变


    # 其余代码保持不变

    """
    5 levels
    0,0,255 ILM-IPL
    0,255,0 IPL-OPL
    0 255 192 OPL-ISOS
    156 255 0 ISOS-RPE
    255 0 0 RPE-BM
    """


"""
提取不同层信息
"""
def color_label_take(label, output_path,folder_name,slice_image_name):
    label=cv2.imread(label)
    # 获取图像形状
    height, width, _ = label.shape

    # 提取颜色信息
    unique_colors = np.unique(label.reshape(-1, label.shape[2]), axis=0)
    unique_colors = unique_colors[unique_colors.any(axis=1)]
    # 确保输出路径存在
    output_path_csv=os.path.join(output_path,'csv',folder_name)
    output_path_image=os.path.join(output_path,'image',folder_name,slice_image_name)
    os.makedirs(output_path_image, exist_ok=True)
    os.makedirs(output_path_csv, exist_ok=True)

    for color in unique_colors:
        # 创建一个与原始图像相同形状的零数组
        result_array_label = np.zeros((height, width, label.shape[2]), dtype=np.uint8)

        # 找出当前颜色的像素位置
        color_mask = (label == color).all(axis=-1)

        # 对于3通道图像，进行连通域分析（需要先转为灰度图）
        gray_color_mask = color_mask.sum(axis=-1) > 0
        _, cc_labels, stats, _ = cv2.connectedComponentsWithStats(gray_color_mask.astype(np.uint8), connectivity=4)

        # 找到面积最大的连通域
        max_area = -1
        max_label = None
        for i, prop in enumerate(regionprops(cc_labels)):
            if prop.area > max_area and prop.area > 0:
                max_area = prop.area
                max_label = i + 1  # 连通域标签从1开始

        # 将最大连通域内的像素恢复成当前颜色
        if max_label is not None:
            result_array_label[color_mask & (cc_labels == max_label)] = color

        # 保存每个颜色的最大连通域结果
        cv2.imwrite(f'{output_path_image}/{str(slice_image_name)}_{tuple(color)}.jpg', result_array_label)

        process_image(f'{output_path_image}/{str(slice_image_name)}_{tuple(color)}.jpg',output_path_csv,color,slice_image_name)
        print(f'oct厚层分析完成---{slice_image_name}')


input_path='/home/imed/apicreate/api_oct-optovue-five_layer_cal/workdir/out_color_data'
output_path='/home/imed/apicreate/api_oct-optovue-five_layer_cal/workdir/OCT_layer_cal/'
os.makedirs(output_path, exist_ok=True)

# 遍历原始图像文件夹
for image_file in os.listdir(input_path):
    if image_file.endswith('.png') or image_file.endswith('.jpg'):
        # 构建完整的输入图像路径
        image_path = os.path.join(input_path, image_file)

        # 定义结果文件夹名称和切片图像名称
        folder_name = 'processed_results'
        slice_image_name = f'processed_{image_file}'

        # 处理当前图像
        color_label_take(image_path, output_path, folder_name, slice_image_name)

print("图像处理完成，并已保存到指定文件夹中。")